# -*- coding: utf-8 -*-
"""
Created on Sat Apr 17 20:17:13 2021

@author: Lenovo-pc
"""


import pandas as pd
import numpy as np
from sklearn import preprocessing


class get_mdp_data_v2:
    def __init__(self,numeric_path,traits_path):
        # self.numeric_path = numeric_path
        # self.traits_path = traits_path
        
        numeric = pd.read_table(numeric_path,header=None) #原始numeric 带有taxa
        numeric = numeric.rename(columns={0:"Taxa"})
        
        traits = pd.read_table(traits_path) #原始traits 带有Taxa
        
        self.traits_flag = traits.columns.tolist()
        self.traits_flag.remove('Taxa')
        #print(self.traits_flag)
        
        #print(self.numeric.columns)
        #self.id = traits['Taxa'] # 获得Taxa
        self.Taxa = numeric['Taxa']
       
        
        # if('Taxa' in self.traits.columns):
        #     self.traits = self.traits.rename(columns={'Taxa':'taxa'})
            
        self.data = pd.merge(traits,numeric,on='Taxa')
        self.data = self.data.dropna()
        self.numeric_cols = numeric.columns.tolist()
        self.numeric_cols.remove('Taxa')
        self.features = numeric[self.numeric_cols]
        
        self.traits_cols = traits.columns.tolist()
        
        
        self.X = self.data[self.numeric_cols]
        self.data_index = self.X.index.tolist()
        self.split()
        
        

    def split(self):

        train_index = self.data_index 
        test_index=[]
       
        index_len = len(self.data_index)
        test_size = int(0.2*index_len)
        i=0
        while(i<test_size):
            randomIndex=int(np.random.uniform(0,index_len))
            if randomIndex in test_index:
                continue
            elif randomIndex not in train_index:
                continue
            else:
                test_index.append(randomIndex) 
                train_index.remove(randomIndex)
                i=i+1
                
        self.train_index = train_index
        self.test_index = test_index
            
    
    def get_traits_flag(self):
        return self.traits_flag
        
        
    def set_traits_flag(self,traits_name):
        self.Y = self.data[traits_name]
    
    def get_split_x(self):
        train_data = self.X.loc[self.train_index]
        test_data = self.X.loc[self.test_index]    
        return train_data.values,test_data.values
        
    def get_one_hot_x(self):
        train_data = self.X.loc[self.train_index]
        test_data = self.X.loc[self.test_index]
    
        categorieslist = [[0,1,2] for i in range(len(self.numeric_cols))]
        
        self.enc = preprocessing.OneHotEncoder(categories=categorieslist)
        self.enc.fit(train_data)
        train_data = self.enc.transform(train_data).toarray()
        test_data = self.enc.transform(test_data).toarray()
        return train_data,test_data
    
    def get_split_y(self):
        train_labels = self.Y.loc[self.train_index]
        test_labels = self.Y.loc[self.test_index]
        return train_labels.values,test_labels.values
    
    def get_all_X(self,is_onehot=False):
        features = self.features
        if is_onehot:
            features = self.enc.transform(features)
            return features

        return features.values
    
    # 保存全部数据的预测值
    def save_into_file(self,predict_X,alg_name,trait_name):
        # self.Taxa
        taxa = self.Taxa.tolist()
        predict = predict_X.tolist()
        assert len(taxa)==len(predict),"ERROR： main  预测数据与taxa长度不一致"
        #save_df = pd.DataFrame([taxa,predict],columns=['taxa',trait_name])
        save_df = pd.DataFrame({'taxa':taxa,trait_name:predict})
        save_path ="./output/{}_{}_result.txt".format(trait_name,alg_name)
        save_df.to_csv(save_path,sep="\t",index=False)
        


if __name__ =="__main__":
    numeric_path ='../data2/mdp_numeric.txt'
    traits_path = '../data2/mdp_Y.txt'
    
    mdp_data = get_mdp_data_v2(numeric_path,traits_path)
    
    train_data,test_data = mdp_data.get_one_hot_x()

    
    
    
    